from functools import partial  # 导入偏函数工具，用于固定函数部分参数
from typing import Callable, List, Optional, Union, Tuple  # 导入类型提示工具
import copy  # 导入复制模块，用于深拷贝对象
import gc  # 导入垃圾回收模块，用于管理内存
import matplotlib.pyplot as plt  # 导入绘图库

import torch  # 导入PyTorch框架
from torchvision.transforms.functional import to_pil_image  # 导入图像转换函数，用于张量转PIL图像
# 导入CLIP相关组件：特征提取器、文本模型、分词器、学习率调度器
from transformers import CLIPFeatureExtractor, CLIPTextModel, CLIPTokenizer, get_cosine_schedule_with_warmup

# 导入Diffusers库中的核心模型：VAE、UNet、安全检查器、调度器
from diffusers.models import AutoencoderKL, UNet2DConditionModel
from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
from diffusers.schedulers import DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler

from torch.optim.lr_scheduler import ReduceLROnPlateau  # 导入学习率调度器（当指标不再改善时降低学习率）

# 导入自定义的修改版稳定扩散管道（基础类）
from src.stable_diffusion.modified_stable_diffusion import ModifiedStableDiffusionPipeline

### 代码参考自：https://github.com/cccntu/efficient-prompt-to-prompt


class InversableStableDiffusionPipeline(ModifiedStableDiffusionPipeline):
    """可逆稳定扩散管道类，继承自修改版稳定扩散管道，支持从图像精确恢复初始噪声（逆过程）"""
    def __init__(self,
        vae,  # VAE模型（变分自编码器，用于图像与潜在空间的转换）
        text_encoder,  # 文本编码器（如CLIP，用于将文本转换为向量）
        tokenizer,  # 分词器（用于文本预处理）
        unet,  # UNet模型（扩散过程的核心，用于预测噪声）
        scheduler,  # 调度器（控制扩散过程的时间步和采样策略）
        safety_checker,  # 安全检查器（用于过滤不当内容）
        feature_extractor,  # 特征提取器（用于安全检查的图像预处理）
        requires_safety_checker: bool = True,  # 是否需要安全检查
    ):
        # 调用父类构造函数初始化组件
        super(InversableStableDiffusionPipeline, self).__init__(vae,
                text_encoder,
                tokenizer,
                unet,
                scheduler,
                safety_checker,
                feature_extractor,
                requires_safety_checker)

    def get_random_latents(self, latents=None, height=512, width=512, generator=None):
        """生成随机初始潜在向量（噪声），作为图像生成的起点"""
        # 确定潜在向量的尺寸（根据UNet和VAE的缩放因子计算）
        height = height or self.unet.config.sample_size * self.vae_scale_factor
        width = width or self.unet.config.sample_size * self.vae_scale_factor

        batch_size = 1  # 批量大小（此处固定为1）
        device = self._execution_device  # 执行设备（GPU/CPU）

        num_channels_latents = self.unet.config.in_channels  # 潜在向量的通道数（通常为4）

        # 准备潜在向量（若未提供则随机生成，符合标准正态分布）
        latents = self.prepare_latents(
            batch_size,
            num_channels_latents,
            height,
            width,
            self.text_encoder.dtype,
            device,
            generator,
            latents,
        )

        return latents  # 返回随机潜在向量

    @torch.inference_mode()  # 禁用梯度计算，加速推理
    def get_text_embedding(self, prompt):
        """将文本提示词编码为向量（文本嵌入）"""
        # 对文本进行分词和编码，转换为模型可处理的输入ID
        text_input_ids = self.tokenizer(
            prompt,
            padding="max_length",  # 填充至最大长度
            truncation=True,  # 截断过长文本
            max_length=self.tokenizer.model_max_length,  # 最大长度（CLIP通常为77）
            return_tensors="pt",  # 返回PyTorch张量
        ).input_ids
        # 使用文本编码器生成文本嵌入（向量表示）
        text_embeddings = self.text_encoder(text_input_ids.to(self.device))[0]
        return text_embeddings  # 返回文本嵌入
    
    @torch.inference_mode()
    def get_image_latents(self, image, sample=True, rng_generator=None):
        """将图像编码为潜在空间向量（通过VAE编码器）"""
        encoding_dist = self.vae.encode(image).latent_dist  # 编码图像得到潜在分布
        if sample:
            encoding = encoding_dist.sample(generator=rng_generator)  # 从分布中采样
        else:
            encoding = encoding_dist.mode()  # 或使用分布的均值（mode）
        latents = encoding * 0.18215  # 应用Stable Diffusion标准缩放因子
        return latents  # 返回图像的潜在向量
    
    @torch.inference_mode()
    def decode_image(self, latents: torch.FloatTensor, **kwargs):
        """将潜在向量解码为图像（通过VAE解码器）"""
        scaled_latents = 1 / 0.18215 * latents  # 逆缩放（抵消编码时的缩放）
        self.vae = self.vae.float()  # 确保VAE为float类型
        # 逐样本解码并拼接结果
        image = [
            self.vae.decode(scaled_latents[i : i + 1]).sample for i in range(len(latents))
        ]
        image = torch.cat(image, dim=0)
        return image  # 返回解码后的图像张量

    def decode_image_for_gradient_float(self, latents: torch.FloatTensor,** kwargs):
        """用于梯度计算的图像解码（深拷贝VAE避免影响原始模型）"""
        scaled_latents = 1 / 0.18215 * latents  # 逆缩放
        vae = copy.deepcopy(self.vae).float()  # 深拷贝VAE，防止梯度传播影响原始模型
        # 逐样本解码并拼接
        image = [
            vae.decode(scaled_latents[i : i + 1]).sample for i in range(len(latents))
        ]
        image = torch.cat(image, dim=0)
        return image  # 返回解码后的图像（用于梯度计算）

    @torch.inference_mode()
    def torch_to_numpy(self, image):
        """将PyTorch图像张量转换为NumPy数组（用于可视化）"""
        image = (image / 2 + 0.5).clamp(0, 1)  # 归一化到[0,1]
        image = image.cpu().permute(0, 2, 3, 1).numpy()  # 调整维度（[B,C,H,W]→[B,H,W,C]）并转移到CPU
        return image  # 返回NumPy格式图像

    def apply_guidance_scale(self, model_output, guidance_scale):
        """应用无分类器引导（classifier-free guidance）"""
        if guidance_scale > 1.0:
            # 若引导尺度>1，分离无条件和有条件噪声预测
            noise_pred_uncond, noise_pred_text = model_output.chunk(2)
            # 结合两者：noise_pred = 无条件预测 + 引导尺度×(有条件预测 - 无条件预测)
            noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
            return noise_pred
        else:
            # 引导尺度=1时，直接返回模型输出（无引导）
            return model_output         

    @torch.inference_mode()
    def forward_diffusion(
        self,
        use_old_emb_i=25,  # Prompt-to-Prompt中切换提示词的时间步索引
        text_embeddings=None,  # 文本嵌入（用于条件生成）
        old_text_embeddings=None,  # 旧提示词嵌入（Prompt-to-Prompt）
        new_text_embeddings=None,  # 新提示词嵌入（Prompt-to-Prompt）
        latents: Optional[torch.FloatTensor] = None,  # 输入潜在向量（图像的潜在表示）
        num_inference_steps: int = 10,  # 推理步数
        guidance_scale: float = 7.5,  # 无分类器引导尺度
        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,  # 回调函数
        callback_steps: Optional[int] = 1,  # 回调函数执行间隔
        inverse_opt=True,  # 是否启用逆过程优化
        inv_order=None,  # 逆过程阶数（1对应一阶算法，2对应二阶算法）
        **kwargs,
    ):  
        """实现扩散过程的逆过程：从图像潜在向量恢复初始噪声（核心方法）"""
        with torch.no_grad():  # 禁用梯度计算（部分步骤可能重新启用）
            # 判断是否使用无分类器引导（引导尺度>1时启用）
            do_classifier_free_guidance = guidance_scale > 1.0

            # 初始化调度器的时间步
            self.scheduler.set_timesteps(num_inference_steps)
            timesteps_tensor = self.scheduler.timesteps.to(self.device)  # 时间步张量（转移到设备）
            latents = latents * self.scheduler.init_noise_sigma  # 初始化潜在向量（乘以初始噪声标准差）

            # 判断是否使用Prompt-to-Prompt（提示词编辑）
            if old_text_embeddings is not None and new_text_embeddings is not None:
                prompt_to_prompt = True
            else:
                prompt_to_prompt = False

            # 若未指定逆过程阶数，使用调度器的求解器阶数
            if inv_order is None:
                inv_order = self.scheduler.solver_order
        
            # 逆过程需反转时间步（从图像向初始噪声推进）
            timesteps_tensor = reversed(timesteps_tensor)

            # 确保模型和张量类型一致（float32）
            self.unet = self.unet.float()
            latents = latents.float()
            text_embeddings = text_embeddings.float()

            # 遍历反转后的时间步，执行逆过程
            for i, t in enumerate(self.progress_bar(timesteps_tensor)):
                # Prompt-to-Prompt：在指定时间步切换新旧提示词嵌入
                if prompt_to_prompt:
                    if i < use_old_emb_i:
                        text_embeddings = old_text_embeddings
                    else:
                        text_embeddings = new_text_embeddings

                # 计算前一个时间步（t的上一步）
                prev_timestep = (
                    t
                    - self.scheduler.config.num_train_timesteps
                    // self.scheduler.num_inference_steps
                )

                # 执行回调函数（若提供）
                if callback is not None and i % callback_steps == 0:
                    callback(i, t, latents)
                

                # 一阶逆过程算法（算法1）：适用于一阶DPM求解器（如DDIM）
                if inv_order < 2 or (inv_order == 2 and i == 0):
                    s = t  # 当前时间步（逆过程中为s）
                    t = prev_timestep  # 前一个时间步（逆过程中为t）
                    
                    # 获取调度器中的关键参数（基于log-SNR和标准差）
                    lambda_s, lambda_t = self.scheduler.lambda_t[s], self.scheduler.lambda_t[t]
                    sigma_s, sigma_t = self.scheduler.sigma_t[s], self.scheduler.sigma_t[t]
                    h = lambda_t - lambda_s  # 时间步间隔（基于log-SNR）
                    alpha_s, alpha_t = self.scheduler.alpha_t[s], self.scheduler.alpha_t[t]
                    phi_1 = torch.expm1(-h)  # 指数项（对应论文中的φ₁函数）

                    # 扩展潜在向量（用于无分类器引导：拼接无条件/有条件输入）
                    latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
                    # 缩放模型输入（根据调度器要求）
                    latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)                
                    
                    # 用UNet预测噪声残差
                    noise_pred = self.unet(latent_model_input, s, encoder_hidden_states=text_embeddings).sample 
                    # 应用无分类器引导
                    noise_pred = self.apply_guidance_scale(noise_pred, guidance_scale)                 
                    # 将噪声预测转换为模型输出（根据调度器规则） q:xt和latents区别
                    model_s = self.scheduler.convert_model_output(noise_pred, t, latents) # 相当于step().prev_sample
                    x_t = latents  # 记录当前潜在向量（用于后续校正）⭐⭐
                    
                    # 算法1第5行：更新潜在向量（初步估计）⭐⭐
                    latents = (sigma_s / sigma_t) * (latents + alpha_t * phi_1 * model_s) # 相当于step().prev_sample后一部分

                    # 算法1第7行：执行固定点校正（优化逆过程精度）
                    if (inverse_opt):
                        # 二阶逆过程的首次迭代使用阶数1校正 ⭐⭐ latents 和 x_t
                        if (inv_order == 2 and i == 0):
                            latents = self.fixedpoint_correction(latents,
                                                                 s, t, x_t, order=1, text_embeddings=text_embeddings, guidance_scale=guidance_scale,
                                                                 step_size=1, scheduler=True)
                        else:
                            latents = self.fixedpoint_correction(latents, s, t, x_t, order=1, text_embeddings=text_embeddings, guidance_scale=guidance_scale,
                                                                 step_size=0.5, scheduler=True)

                # 二阶逆过程算法（算法2）：适用于高阶DPM求解器（如DPM-Solver++(2M)）
                elif inv_order == 2:
                    with torch.no_grad():
                        # 处理非最后一个时间步
                        if (i + 1 < len(timesteps_tensor)):                          
                            y = latents.clone()  # 复制当前潜在向量（用于高阶项近似）

                            s = t  # 当前时间步
                            t = prev_timestep  # 前一个时间步
                            r = timesteps_tensor[i + 1] if i+1 < len(timesteps_tensor) else 0  # 下一个时间步（用于高阶项）
                            
                            # 算法2第3-6行：使用细粒度朴素DDIM逆过程估计历史潜在向量（黄线部分）
                            for tt in range(t,s,10):  # 时间步间隔10，细粒度采样
                                ss = tt + 10  # 下一个细粒度时间步
                                # 获取调度器参数
                                lambda_s, lambda_t = self.scheduler.lambda_t[ss], self.scheduler.lambda_t[tt]
                                sigma_s, sigma_t = self.scheduler.sigma_t[ss], self.scheduler.sigma_t[tt]
                                h = lambda_t - lambda_s
                                alpha_s, alpha_t = self.scheduler.alpha_t[ss], self.scheduler.alpha_t[tt]
                                phi_1 = torch.expm1(-h)

                                # 准备模型输入（带引导）
                                y_input = torch.cat([y] * 2) if do_classifier_free_guidance else y
                                y_input = self.scheduler.scale_model_input(y_input, tt)

                                # 预测噪声并更新潜在向量（朴素DDIM逆过程）
                                noise_pred = self.unet(y_input, ss, encoder_hidden_states=text_embeddings).sample
                                noise_pred = self.apply_guidance_scale(noise_pred, guidance_scale)    
                                model_s = self.scheduler.convert_model_output(noise_pred, tt, y)
                                y = (sigma_s / sigma_t) * (y + alpha_t * phi_1 * model_s)  # 对应算法1第5行
                            y_t = y.clone()  # 保存估计的历史潜在向量
                            
                            # 继续估计更早的潜在向量（用于高阶项）
                            for tt in range(s, r,10):
                                ss = tt + 10
                                lambda_s, lambda_t = self.scheduler.lambda_t[ss], self.scheduler.lambda_t[tt]
                                sigma_s, sigma_t = self.scheduler.sigma_t[ss], self.scheduler.sigma_t[tt]
                                h = lambda_t - lambda_s
                                alpha_s, alpha_t = self.scheduler.alpha_t[ss], self.scheduler.alpha_t[tt]
                                phi_1 = torch.expm1(-h)

                                y_input = torch.cat([y] * 2) if do_classifier_free_guidance else y
                                y_input = self.scheduler.scale_model_input(y_input, tt)

                                model_s = self.unet(y_input, ss, encoder_hidden_states=text_embeddings).sample
                                noise_pred = self.apply_guidance_scale(model_s, guidance_scale)    
                                model_s = self.scheduler.convert_model_output(noise_pred, tt, y) 
                                y = (sigma_s / sigma_t) * (y + alpha_t * phi_1 * model_s)


                            # 算法2第8-12行：使用向后欧拉方法结合高阶项近似更新潜在向量（蓝线部分）
                            t = prev_timestep
                            s = timesteps_tensor[i]
                            r = timesteps_tensor[i+1]
                            
                            # 获取调度器参数
                            lambda_s, lambda_t = self.scheduler.lambda_t[s], self.scheduler.lambda_t[t]
                            sigma_s, sigma_t = self.scheduler.sigma_t[s], self.scheduler.sigma_t[t]
                            h = lambda_t - lambda_s
                            alpha_s, alpha_t = self.scheduler.alpha_t[s], self.scheduler.alpha_t[t]
                            phi_1 = torch.expm1(-h)
                            
                            x_t = latents  # 记录当前潜在向量
                            
                            # 估计高阶项所需的模型输出
                            y_t_model_input = torch.cat([y_t] * 2) if do_classifier_free_guidance else y_t
                            y_t_model_input = self.scheduler.scale_model_input(y_t_model_input, s)
                            
                            noise_pred = self.unet(y_t_model_input, s, encoder_hidden_states=text_embeddings).sample 
                            noise_pred = self.apply_guidance_scale(noise_pred, guidance_scale)
                            model_s_output = self.scheduler.convert_model_output(noise_pred, s, y_t)
                            
                            y_model_input = torch.cat([y] * 2) if do_classifier_free_guidance else y
                            y_model_input = self.scheduler.scale_model_input(y_model_input, r)
                            
                            noise_pred = self.unet(y_model_input, r, encoder_hidden_states=text_embeddings).sample
                            noise_pred = self.apply_guidance_scale(noise_pred, guidance_scale)                           
                            model_r_output = self.scheduler.convert_model_output(noise_pred, r, y)
                            
                            latents = y_t.clone()  # 算法2第7行：初始化潜在向量为估计的历史值
                            
                            # 算法2第11行：执行固定点校正（带高阶项近似）
                            if inverse_opt:
                                latents = self.fixedpoint_correction(latents, s, t, x_t, order=2, r=r,
                                                                    model_s_output=model_s_output, model_r_output=model_r_output, text_embeddings=text_embeddings, guidance_scale=guidance_scale,
                                                                    step_size=10/t, scheduler=False) 
                            
                        # 处理最后一个时间步
                        elif (i + 1 == len(timesteps_tensor)):
                            s = t
                            t = prev_timestep
                            
                            # 获取调度器参数
                            lambda_s, lambda_t = self.scheduler.lambda_t[s], self.scheduler.lambda_t[t]
                            sigma_s, sigma_t = self.scheduler.sigma_t[s], self.scheduler.sigma_t[t]
                            h = lambda_t - lambda_s
                            alpha_s, alpha_t = self.scheduler.alpha_t[s], self.scheduler.alpha_t[t]
                            phi_1 = torch.expm1(-h)

                            # 准备模型输入
                            latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents                          
                            latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
                            
                            # 预测噪声并更新潜在向量
                            noise_pred = self.unet(latent_model_input, s, encoder_hidden_states=text_embeddings).sample
                            noise_pred = self.apply_guidance_scale(noise_pred, guidance_scale)        
                            model_s = self.scheduler.convert_model_output(noise_pred, t, latents)

                            x_t = latents  # 记录当前潜在向量
                            
                            # 算法2第16行：更新潜在向量
                            latents = (sigma_s / sigma_t) * (latents + alpha_t * phi_1 * model_s)
                            
                            # 算法2第17行：执行固定点校正
                            if (inverse_opt):
                                latents = self.fixedpoint_correction(latents, s, t, x_t, order=1, text_embeddings=text_embeddings, guidance_scale=guidance_scale,
                                                                     step_size=10/t, scheduler=True)   
                        else:
                            raise Exception("Index Error!")  # 索引错误处理
                else:
                    pass  # 未实现更高阶的逆过程

        return latents  # 返回恢复的初始噪声（潜在向量）

    @torch.inference_mode()
    def fixedpoint_correction(self, x, s, t, x_t, r=None, order=1, n_iter=500, step_size=0.1, th=1e-3, 
                                model_s_output=None, model_r_output=None, text_embeddings=None, guidance_scale=3.0, 
                                scheduler=False, factor=0.5, patience=20, anchor=False, warmup=True, warmup_time=20):
        """
        固定点校正（Fixed-point Correction）
        ------------------------------------
        通过迭代优化潜在向量 latent，保证反演过程的精确性。
        一阶方法依赖 (s -> t)，二阶方法额外引入 (r -> s) 的高阶项近似。

        参数
        ----------
        x : torch.Tensor
            当前时间步 s 的潜在向量估计值（latent），待优化的对象。初始值通常来自反演推导公式。
        s : int
            当前时间步索引 (source step)。
        t : int
            下一目标时间步索引 (target step)。
        x_t : torch.Tensor
            在目标时间步 t 的真实潜在向量，用于计算校正误差。
        r : int, optional
            s 的上一个时间步索引，仅在二阶校正 (order=2) 时需要。
        order : int, default=1
            固定点校正的阶数。
            - 1: 一阶方法（s → t，常规校正）
            - 2: 二阶方法（r → s → t，高阶校正）
        n_iter : int, default=500
            最大迭代次数。超过此次数即停止。
        step_size : float, default=0.1
            每次更新潜在向量的步长，相当于学习率。
        th : float, default=1e-3
            收敛阈值。当 MSE 损失小于该值时提前停止。
        model_s_output : torch.Tensor, optional
            在时间步 s 上的 U-Net 输出（噪声预测经调度器转换）。
            二阶校正时需要。
        model_r_output : torch.Tensor, optional
            在时间步 r 上的 U-Net 输出。二阶校正时需要。
        text_embeddings : torch.Tensor, optional
            条件文本嵌入（来自文本/音频编码器），作为 U-Net 输入。
        guidance_scale : float, default=3.0
            无分类器指导 (CFG) 系数。
            - 1.0: 不使用指导
            - >1.0: 放大条件与无条件预测差异
        scheduler : bool, default=False
            是否启用动态步长调度器 (StepScheduler)。
        factor : float, default=0.5
            步长调度器的衰减因子。
        patience : int, default=20
            步长调度器的容忍度，若损失不改善超过此步数则调整。
        anchor : bool, default=False
            是否启用锚定机制（仅 order=2 使用），
            在迭代中部分回拉初始值以增强稳定性。
        warmup : bool, default=True
            是否启用步长预热。在前 warmup_time 步内逐渐增大步长。
        warmup_time : int, default=20
            步长预热的迭代次数。

        返回
        -------
        torch.Tensor
            优化后的潜在向量（修正后的 latent）。
        """
        do_classifier_free_guidance = guidance_scale > 1.0  # 是否启用无分类器引导
        if order==1:
            """一阶固定点校正（对应算法1）"""
            input = x.clone()  # 待优化的潜在向量（初始化为当前估计值）
            original_step_size = step_size  # 保存初始步长（用于预热）
            
            # 初始化步长调度器（当损失不再改善时降低步长）
            if scheduler:
                step_scheduler = StepScheduler(current_lr=step_size, factor=factor, patience=patience)

            # 获取调度器参数
            lambda_s, lambda_t = self.scheduler.lambda_t[s], self.scheduler.lambda_t[t]
            alpha_s, alpha_t = self.scheduler.alpha_t[s], self.scheduler.alpha_t[t]
            sigma_s, sigma_t = self.scheduler.sigma_t[s], self.scheduler.sigma_t[t]
            h = lambda_t - lambda_s
            phi_1 = torch.expm1(-h)

            # 迭代优化（最多n_iter次）
            for i in range(n_iter):
                # 步长预热：前warmup_time步逐渐增加步长
                if warmup:
                    if i < warmup_time:
                        step_size = original_step_size * (i+1)/(warmup_time)
                
                # 准备模型输入（带引导）
                latent_model_input = (torch.cat([input] * 2) if do_classifier_free_guidance else input)
                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
                
                # 预测噪声并转换为模型输出⭐⭐⭐⭐
                noise_pred = self.unet(latent_model_input , s, encoder_hidden_states=text_embeddings).sample
                noise_pred = self.apply_guidance_scale(noise_pred, guidance_scale)                   
                model_output = self.scheduler.convert_model_output(noise_pred, s, input)
                # 计算预测的当前时间步潜在向量（用于与真实值比较） 相当于step().prev_sample
                x_t_pred = (sigma_t / sigma_s) * input - (alpha_t * phi_1 ) * model_output

                # 计算损失（MSE：预测值与真实值的差异）
                loss = torch.nn.functional.mse_loss(x_t_pred, x_t, reduction='sum')
                
                # 若损失小于阈值，提前停止迭代
                if loss.item() < th:
                    break                
                
                # 前向步长法更新：input = input - 步长×(预测误差)
                input = input - step_size * (x_t_pred - x_t)

                # 更新步长（若启用调度器）
                if scheduler:
                    step_size = step_scheduler.step(loss)

            return input  # 返回优化后的潜在向量        
        
        elif order==2:
            """二阶固定点校正（对应算法2，带高阶项近似）"""
            assert r is not None  # 确保提供了下一个时间步r（用于高阶项）
            input = x.clone()  # 待优化的潜在向量
            original_step_size = step_size  # 初始步长
            
            # 初始化步长调度器
            if scheduler:
                step_scheduler = StepScheduler(current_lr=step_size, factor=factor, patience=patience)
            
            # 获取调度器参数（含高阶项所需的r）
            lambda_r, lambda_s, lambda_t = self.scheduler.lambda_t[r], self.scheduler.lambda_t[s], self.scheduler.lambda_t[t]
            sigma_r, sigma_s, sigma_t = self.scheduler.sigma_t[r], self.scheduler.sigma_t[s], self.scheduler.sigma_t[t]
            alpha_s, alpha_t = self.scheduler.alpha_t[s], self.scheduler.alpha_t[t]
            h_0 = lambda_s - lambda_r  # 时间步间隔（r到s）
            h = lambda_t - lambda_s  # 时间步间隔（s到t）
            r0 = h_0 / h  # 间隔比例（用于高阶项）
            phi_1 = torch.expm1(-h)  # 指数项
            
            # 迭代优化
            for i in range(n_iter):
                # 步长预热
                if warmup:
                    if i < warmup_time:
                        step_size = original_step_size * (i+1)/(warmup_time)

                # 准备模型输入（带引导）
                latent_model_input = torch.cat([input] * 2) if do_classifier_free_guidance else input
                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)

                # 预测噪声并转换为模型输出
                noise_pred = self.unet(latent_model_input, s, encoder_hidden_states=text_embeddings).sample
                noise_pred = self.apply_guidance_scale(noise_pred, guidance_scale) 
                model_output = self.scheduler.convert_model_output(noise_pred, s, input)
                
                # 计算预测的当前时间步潜在向量
                x_t_pred = (sigma_t / sigma_s) * input - (alpha_t * phi_1) * model_output
                
                # 高阶项近似（仅在首次迭代计算，固定为常数）
                if i==0:
                    d = (1./ r0) * (model_s_output - model_r_output)  # 高阶项导数近似
                    diff_term = 0.5 * alpha_t * phi_1 * d  # 高阶校正项

                # 应用高阶项校正
                x_t_pred = x_t_pred - diff_term
                
                # 计算损失
                loss = torch.nn.functional.mse_loss(x_t_pred, x_t, reduction='sum')

                # 若损失小于阈值，提前停止
                if loss.item() < th:
                    break                

                # 前向步长法更新
                input = input - step_size * (x_t_pred - x_t)

                # 更新步长（若启用调度器）
                if scheduler:
                    step_size = step_scheduler.step(loss)
                # 锚定机制：缓慢向初始估计值靠近，增强稳定性
                if anchor:
                    input = (1 - 1/(i+2)) * input + (1/(i+2))*x
            return input  # 返回优化后的潜在向量（带高阶校正）
        else:
            raise NotImplementedError  # 未实现更高阶校正

    def decoder_inv(self, x):
        """解码器逆过程：通过优化找到使VAE解码结果与输入图像匹配的潜在向量（减少重建误差）"""
        """
        输入：
            x : 图像数据（形状为(1, 3, 512, 512)）
        输出：
            z : 优化后的潜在向量（形状为(1, 4, 64, 64)）
        目标：最小化 ||VAE解码(z) - x||（即让潜在向量z解码后尽可能接近输入图像x）
        """
        input = x.clone().float()  # 输入图像（深拷贝避免修改原始数据）

        # 初始化潜在向量（使用VAE编码器的结果作为初始值）
        z = self.get_image_latents(x).clone().float()
        z.requires_grad_(True)  # 启用梯度（需优化z）

        # 定义损失函数（MSE）和优化器
        loss_function = torch.nn.MSELoss(reduction='sum')
        optimizer = torch.optim.Adam([z], lr=0.1)  # Adam优化器，初始学习率0.1
        # 学习率调度器（余弦退火+预热）
        lr_scheduler = get_cosine_schedule_with_warmup(optimizer, num_warmup_steps=10, num_training_steps=100)

        # 迭代优化（100次）
        for i in self.progress_bar(range(100)):
            x_pred = self.decode_image_for_gradient_float(z)  # 解码潜在向量得到预测图像

            loss = loss_function(x_pred, input)  # 计算预测图像与输入图像的MSE损失
            
            # 反向传播与参数更新
            optimizer.zero_grad()  # 清空梯度
            loss.backward()  # 计算梯度
            optimizer.step()  # 更新潜在向量z
            lr_scheduler.step()  # 更新学习率
        return z  # 返回优化后的潜在向量
    
   
        
class StepScheduler(ReduceLROnPlateau):
    """步长调度器：当损失不再改善时降低步长（继承自PyTorch的ReduceLROnPlateau）"""
    def __init__(self, mode='min', current_lr=0, factor=0.1, patience=10,
                 threshold=1e-4, threshold_mode='rel', cooldown=0,
                 min_lr=0, eps=1e-8, verbose=False):
        if factor >= 1.0:
            raise ValueError('衰减因子必须小于1.0。')
        self.factor = factor  # 学习率衰减因子
        if current_lr == 0:
            raise ValueError('初始步长不能为0。')

        self.min_lr = min_lr  # 最小学习率
        self.current_lr = current_lr  # 当前学习率
        self.patience = patience  # 容忍损失不变的迭代次数
        self.verbose = verbose  # 是否打印日志
        self.cooldown = cooldown  # 冷却期（衰减后暂停衰减的迭代次数）
        self.cooldown_counter = 0  # 冷却期计数器
        self.mode = mode  # 优化模式（'min'表示最小化损失，'max'表示最大化指标）
        self.threshold = threshold  # 损失变化的阈值（小于此值视为无改善）
        self.threshold_mode = threshold_mode  # 阈值模式（'rel'相对变化，'abs'绝对变化）
        self.best = None  # 最佳损失值
        self.num_bad_epochs = None  # 连续损失无改善的次数
        self.mode_worse = None  # 模式对应的"更差"值（min模式为inf，max模式为-inf）
        self.eps = eps  # 学习率最小变化量
        self.last_epoch = 0  # 最后迭代次数
        # 初始化"是否更好"的判断函数
        self._init_is_better(mode=mode, threshold=threshold,
                             threshold_mode=threshold_mode)
        self._reset()  # 重置状态

    def step(self, metrics, epoch=None):
        """根据当前损失更新学习率"""
        current = float(metrics)  # 将损失转换为浮点数
        if epoch is None:
            epoch = self.last_epoch + 1
        else:
            import warnings
            warnings.warn("EPOCH_DEPRECATION_WARNING", UserWarning)
        self.last_epoch = epoch

        # 判断当前损失是否优于最佳损失
        if self.is_better(current, self.best):
            self.best = current
            self.num_bad_epochs = 0
        else:
            self.num_bad_epochs += 1  # 连续无改善次数+1

        # 冷却期处理
        if self.in_cooldown:
            self.cooldown_counter -= 1
            self.num_bad_epochs = 0  # 冷却期内不计数

        # 若连续无改善次数超过容忍值，衰减学习率
        if self.num_bad_epochs > self.patience:
            self._reduce_lr(epoch)
            self.cooldown_counter = self.cooldown  # 进入冷却期
            self.num_bad_epochs = 0

        return self.current_lr  # 返回当前学习率

    def _reduce_lr(self, epoch):
        """衰减学习率"""
        old_lr = self.current_lr
        # 新学习率 = 当前学习率×衰减因子（不低于最小学习率）
        new_lr = max(self.current_lr * self.factor, self.min_lr)
        # 若学习率变化超过最小阈值，更新学习率
        if old_lr - new_lr > self.eps:
            self.current_lr = new_lr
            if self.verbose:
                epoch_str = ("%.2f" if isinstance(epoch, float) else
                            "%.5d") % epoch
                print(f'迭代 {epoch_str}: 学习率衰减至 {new_lr:.4e}。')